#!d:\anaconda3\envs\office_automation\python.exe -u
# -*- coding=utf-8 -*-
r"""
author: git config IVEN_CN && git config 13377529851@QQ.com
Date: 2024-08-22 09:23:11 +0800
LastEditTime: 2024-08-22 16:07:20 +0800
FilePath: \office_automation\color_div\color_divUI.py
details: 用于颜色识别的UI界面

Copyright (c) 2024 by IVEN, All Rights Reserved. 
"""
import os
import sys
import PIL.Image
import PyQt5.QtWidgets as qw
import numpy as np
try:
    from color_divider import color_detect, delete_empty_folders, get_file_path
except ImportError:
    from color_div.color_divider import color_detect, delete_empty_folders, get_file_path

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from cutter.cut import cut_cut, resize1200


class SingleColorDivWindow(qw.QWidget):
    """单颜色识别窗口
    ----
    * 采用半自动的识别方式, 需要用户手动把指定颜色的图片放入颜色文件夹
    * 全自动的识别分类精准度会受到背景颜色的影响, 这里尽量避免使用全自动的方式"""

    COLOR_LST = [
        "白色2100",
        "黑色2010",
        "杏色2170",
        "粉红2340",
        "红色2200",
        "黄色2400",
        "深灰2040",
        "卡其2970",
        "克莱因蓝2500",
        "彩兰2670",
        "花灰2080",
        "灰蓝2690",
    ]

    def __init__(self) -> None:
        super().__init__()
        self.setWindowTitle("单颜色分类")
        self.main_layout = qw.QHBoxLayout(self)
        self.left_layout = qw.QVBoxLayout()
        self.right_layout = qw.QVBoxLayout()
        self.main_layout.addLayout(self.left_layout)
        self.main_layout.addLayout(self.right_layout)
        self.setLayout(self.main_layout)

        # region 添加选择路径的按钮和标签
        path_layout = qw.QHBoxLayout()
        self.path_button = qw.QPushButton("选择路径")
        self.path_button.clicked.connect(self.select_path)
        path_layout.addWidget(self.path_button)

        self.path_label = qw.QLabel("未选择路径")
        path_layout.addWidget(self.path_label)

        self.left_layout.addLayout(path_layout)
        # endregion

        # region 颜色复选框
        color_box = qw.QGroupBox("颜色")
        color_layout = qw.QVBoxLayout()
        color_box.setLayout(color_layout)

        self.checkboxes = []
        for color in self.COLOR_LST:
            checkbox = qw.QCheckBox(color)
            self.checkboxes.append(checkbox)
            color_layout.addWidget(checkbox)

        self.left_layout.addWidget(color_box)
        # endregion

        # region 新建颜色文件夹按钮
        self.mkdir_button = qw.QPushButton("新建颜色文件夹")
        self.mkdir_button.setSizePolicy(
            qw.QSizePolicy.Expanding, qw.QSizePolicy.Expanding
        )
        self.mkdir_button.clicked.connect(self.mkdir_callback)
        self.left_layout.addWidget(self.mkdir_button)
        # endregion

        # region 删除空文件按钮
        self.delete_button = qw.QPushButton("删除空文件夹")
        self.delete_button.setSizePolicy(
            qw.QSizePolicy.Expanding, qw.QSizePolicy.Expanding
        )
        self.delete_button.clicked.connect(self.clean_empty_folders)
        self.left_layout.addWidget(self.delete_button)
        # endregion

        # region 文件复选框
        self.file_box = qw.QGroupBox("图片文件")
        self.file_layout = qw.QVBoxLayout()
        self.file_box.setLayout(self.file_layout)

        self.path = self.path_label.text()
        self.update_file_checkboxes()

        self.right_layout.addWidget(self.file_box)
        # endregion

        # region 裁剪按钮
        cut_button = qw.QPushButton("裁剪")
        cut_button.setSizePolicy(qw.QSizePolicy.Expanding, qw.QSizePolicy.Expanding)
        cut_button.clicked.connect(self.cut_img)
        self.right_layout.addWidget(cut_button)

        # endregion

    def update_file_checkboxes(self):
        files = get_file_path(self.path)
        filenames = [os.path.basename(i) for i in files]
        filenames = [i for i in filenames if i.endswith(".jpg")]

        # 清空原有的复选框
        for i in range(self.file_layout.count()):
            self.file_layout.itemAt(i).widget().deleteLater()  # type:ignore

        self.check_files = []
        for filename in filenames:
            checkbox = qw.QCheckBox(filename)
            self.check_files.append(checkbox)
            self.file_layout.addWidget(checkbox)

    def mkdir_callback(self):
        path = self.path_label.text()
        selected_colors = [cb.text() for cb in self.checkboxes if cb.isChecked()]
        for i in selected_colors:
            try:
                os.mkdir(os.path.join(path, i))
            except FileExistsError:
                pass

    def clean_empty_folders(self):
        path = self.path_label.text()
        delete_empty_folders(path)

    def select_path(self):
        options = qw.QFileDialog.Options()
        options |= qw.QFileDialog.ShowDirsOnly
        directory = qw.QFileDialog.getExistingDirectory(self, "选择文件夹", "", options)
        if directory:
            self.path_label.setText(directory)
            self.path = directory
        self.update_file_checkboxes()

    def cut_img(self):
        selected_files = [cb.text() for cb in self.check_files if cb.isChecked()]
        selected_files = tuple(selected_files)
        resize1200(*selected_files, path=self.path)
        cut_cut(*selected_files, path=self.path)


# XXX: 未测试, 可能有问题
class MoreColorDivWindow(qw.QWidget):
    COLOR_LST = [
        "白色2100",
        "黑色2010",
        "杏色2170",
        "粉红2340",
        "红色2200",
        "黄色2400",
        "深灰2040",
        "卡其2970",
        "克莱因蓝2500",
        "彩兰2670",
        "花灰2080",
        "灰蓝2690",
    ]

    def __init__(self) -> None:
        super().__init__()
        self.setWindowTitle("多颜色分类")
        self.main_layout = qw.QHBoxLayout(self)
        self.left_layout = qw.QVBoxLayout()
        self.right_layout = qw.QVBoxLayout()
        self.main_layout.addLayout(self.left_layout)
        self.main_layout.addLayout(self.right_layout)
        self.setLayout(self.main_layout)

        # region 添加选择路径的按钮和标签
        path_layout = qw.QHBoxLayout()
        self.path_button = qw.QPushButton("选择路径")
        self.path_button.clicked.connect(self.select_path)
        path_layout.addWidget(self.path_button)

        self.path_label = qw.QLabel("未选择路径")
        path_layout.addWidget(self.path_label)

        self.left_layout.addLayout(path_layout)
        # endregion

        # region 颜色复选框
        color_box = qw.QGroupBox("颜色")
        color_layout = qw.QVBoxLayout()
        color_box.setLayout(color_layout)

        self.checkboxes = []
        for color in self.COLOR_LST:
            checkbox = qw.QCheckBox(color)
            self.checkboxes.append(checkbox)
            color_layout.addWidget(checkbox)

        self.left_layout.addWidget(color_box)
        # endregion

        # region 文件复选框
        self.file_box = qw.QGroupBox("图片文件")
        self.file_layout = qw.QVBoxLayout()
        self.file_box.setLayout(self.file_layout)

        self.path = self.path_label.text()
        SingleColorDivWindow.update_file_checkboxes(self)

        self.right_layout.addWidget(self.file_box)
        # endregion

        # region 裁剪按钮
        cut_button = qw.QPushButton("裁剪")
        cut_button.setSizePolicy(qw.QSizePolicy.Expanding, qw.QSizePolicy.Expanding)
        cut_button.clicked.connect(self.cut_img)
        self.right_layout.addWidget(cut_button)
        # endregion

        # region 一键分类按钮
        auto_button = qw.QPushButton("一键分类")
        auto_button.setSizePolicy(qw.QSizePolicy.Expanding, qw.QSizePolicy.Expanding)
        auto_button.clicked.connect(self.auto_div)
        self.right_layout.addWidget(auto_button)
        # endregion

    def select_path(self):
        options = qw.QFileDialog.Options()
        options |= qw.QFileDialog.ShowDirsOnly
        directory = qw.QFileDialog.getExistingDirectory(self, "选择文件夹", "", options)
        if directory:
            self.path_label.setText(directory)
            self.path = directory
        self.update_file_checkboxes()

    def cut_img(self):
        selected_files = [cb.text() for cb in self.check_files if cb.isChecked()]
        selected_files = tuple(selected_files)
        resize1200(*selected_files, path=self.path)
        cut_cut(*selected_files, path=self.path)

    def update_file_checkboxes(self):
        files = get_file_path(self.path)
        filenames = [os.path.basename(i) for i in files]
        filenames = [i for i in filenames if i.endswith(".jpg")]

        # 清空原有的复选框
        for i in range(self.file_layout.count()):
            self.file_layout.itemAt(i).widget().deleteLater()  # type:ignore

        self.check_files = []
        for filename in filenames:
            checkbox = qw.QCheckBox(filename)
            self.check_files.append(checkbox)
            self.file_layout.addWidget(checkbox)

    def auto_div(self):
        # 配置文件路径
        cfg_paths = os.listdir(f"{self.path}/..")
        cfg_names = [i for i in cfg_paths if i.endswith(".npy")]  # 只有npy文件, 文件名
        color_cfg_names = [
            i for i in cfg_names if "area" not in i
        ]  # 只有颜色npy文件,  文件名
        color_names = [
            i.split(".")[0] for i in color_cfg_names
        ]  # 只有颜色名字  ->  文件夹名

        # 创建文件夹
        for i in color_names:
            try:
                os.mkdir(os.path.join(self.path, i))
            except FileExistsError:
                pass

        # 多颜色识别
        dict_ = {}
        path = get_file_path(self.path)     # 图片路径列表
        for file_path in path:
            img = self.imread(file_path)
            if img is None:
                continue

            dict_[file_path] = []

            for cfg in cfg_names:
                # self.path父级文件夹下的配置文件
                color_cfg_path = os.path.join(self.path, "..", cfg)
                area_cfg_path = os.path.join(self.path, "..", "area_cloth.npy")
                if color_detect(img, color_cfg_path, area_cfg_path) and len(dict_[file_path]) < 3:
                    dict_[file_path].append(cfg.split(".")[0])

            # 文件夹名
            class_name: str = ""
            class_name = "+".join(dict_[file_path])

            # 创建文件夹
            try:
                os.mkdir(os.path.join(self.path, class_name))
            except FileExistsError:
                pass

            # 移动文件
            os.system(f"move {i} {os.path.join(self.path, class_name)}")

    def imread(self, path: str):
        """
        读取图片

        :param path: 图片路径
        """
        img = PIL.Image.open(path)
        img = np.asarray(img)
        return img


if __name__ == "__main__":
    app = qw.QApplication(sys.argv)
    window = MoreColorDivWindow()
    window.show()
    sys.exit(app.exec_())
